#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <string.h>
#include <errno.h>
#include <linux/vfio.h>

#define VFIO_DEV "/dev/vfio/vfio"
#define CONTAINER_IOMMU_GROUP "/dev/vfio/5"  // 根据您的iommu_group=5设置
#define BUFFER_SIZE (4 * 1024 * 1024)       // 4MB缓冲区
#define _GNU_SOURCE                         // 用于posix_memalign

int main() {
    // 1. 打开VFIO设备
    int container_fd = open(VFIO_DEV, O_RDWR);
    if (container_fd < 0) {
        perror("无法打开VFIO设备");
        goto release;
    }
    
    // 2. 检查VFIO API版本
    int vfio_api_version = ioctl(container_fd, VFIO_GET_API_VERSION);
    if (vfio_api_version != VFIO_API_VERSION) {
        printf("The VFIO API version{%d} does not match VFIO_API_VERSION{%d}.",
            vfio_api_version, VFIO_API_VERSION);
        goto release;
    }
    if (!ioctl(container_fd, VFIO_CHECK_EXTENSION, VFIO_TYPE1_IOMMU)) {
        printf("Doesn't support the IOMMU driver we want.\n");
        goto release;
    }
    
    // 3. 打开IOMMU组
    int group_fd = open(CONTAINER_IOMMU_GROUP, O_RDWR);
    if (group_fd < 0) {
        printf("Open %s failed, %s\n", CONTAINER_IOMMU_GROUP, strerror(errno));
        goto release;
    } else {
        printf("Open %s success.\n", CONTAINER_IOMMU_GROUP);
    }
    
    // 4. 检查组状态
    struct vfio_group_status group_status = { .argsz = sizeof(group_status) };
    if (ioctl(group_fd, VFIO_GROUP_GET_STATUS, &group_status)) {
        printf("Get group status failed, %s\n", strerror(errno));
        goto release;
    }
    
    if (!(group_status.flags & VFIO_GROUP_FLAGS_VIABLE)) {
        printf("IOMMU group{%s} is not viable.\n", CONTAINER_IOMMU_GROUP);
        goto release;
    }
    
    // 5. 将组加入容器
    if (ioctl(group_fd, VFIO_GROUP_SET_CONTAINER, &container_fd)) {
        printf("Set container failed, group{%s} %s\n", CONTAINER_IOMMU_GROUP, strerror(errno));
        goto release;
    }

    if (ioctl(container_fd, VFIO_SET_IOMMU, VFIO_TYPE1_IOMMU)) {
        printf("Set IOMMU failed, %s\n", strerror(errno));
        goto release;
    }

    struct vfio_iommu_type1_info iommu_info = { .argsz = sizeof(iommu_info) };
    if (ioctl(container_fd, VFIO_IOMMU_GET_INFO, &iommu_info)) {
        printf("Get IOMMU info failed, %s\n", strerror(errno));
        goto release;
    }

    struct vfio_iommu_type1_dma_map dma_map = { .argsz = sizeof(dma_map) };
    dma_map.vaddr = (__u64)mmap(0, 1024 * 1024, PROT_READ | PROT_WRITE,
        MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
	dma_map.size = 1024 * 1024;
	dma_map.iova = 0; /* 1MB starting at 0x0 from device view */
	dma_map.flags = VFIO_DMA_MAP_FLAG_READ | VFIO_DMA_MAP_FLAG_WRITE;
    if (ioctl(container_fd, VFIO_IOMMU_MAP_DMA, &dma_map)) {
        printf("Map DMA failed, %s\n", strerror(errno));
        goto release;
    }

    int device_fd = ioctl(group_fd, VFIO_GROUP_GET_DEVICE_FD, "0000:00:17.0");
    if (device_fd < 0) {
        printf("Get device fd failed, %s\n", strerror(errno));
        goto release;
    }

    struct vfio_device_info device_info = { .argsz = sizeof(device_info) };
    if (ioctl(device_fd, VFIO_DEVICE_GET_INFO, &device_info)) {
        printf("Get device info failed, %s\n", strerror(errno));
        goto release;
    }
    printf("device_info.num_regions{%u} num_irqs{%u} flags{%x}\n",
        device_info.num_regions, device_info.num_irqs, device_info.flags);
    if (!(device_info.flags & VFIO_DEVICE_FLAGS_PCI)) {
        printf("Not a PCI device.\n");
        goto release;
    }
    for (int i = 0; i < device_info.num_regions && i <= VFIO_PCI_CONFIG_REGION_INDEX; i++) {
        struct vfio_region_info tmp = { .argsz = sizeof(tmp), .index = i };
        if (ioctl(device_fd, VFIO_DEVICE_GET_REGION_INFO, &tmp)) {
            printf("Failed to get region info {%s}\n", strerror(errno));
            goto release;
        }
        printf("region_info index{%u} size{%lu} offset{%lu} READ{%d} WRITE{%d} MMAP{%d} CAPS{%d}\n",
            tmp.index, tmp.size, tmp.offset,
            tmp.flags & VFIO_REGION_INFO_FLAG_READ ? 1:0,
            tmp.flags & VFIO_REGION_INFO_FLAG_WRITE ? 1:0,
            tmp.flags & VFIO_REGION_INFO_FLAG_MMAP ? 1:0,
            tmp.flags & VFIO_REGION_INFO_FLAG_CAPS ? 1:0);
    }
    // 映射配置空间
    struct vfio_region_info reg_info_config = { .argsz = sizeof(reg_info_config), .index = VFIO_PCI_CONFIG_REGION_INDEX };
    if (ioctl(device_fd, VFIO_DEVICE_GET_REGION_INFO, &reg_info_config)) {
        printf("Failed to get region info {%s}\n", strerror(errno));
        goto release;
    }
    printf("region_info index{%u} size{%lu} offset{%lu} READ{%d} WRITE{%d} MMAP{%d} CAPS{%d}\n",
        reg_info_config.index, reg_info_config.size, reg_info_config.offset,
        reg_info_config.flags & VFIO_REGION_INFO_FLAG_READ,
        reg_info_config.flags & VFIO_REGION_INFO_FLAG_WRITE,
        reg_info_config.flags & VFIO_REGION_INFO_FLAG_MMAP,
        reg_info_config.flags & VFIO_REGION_INFO_FLAG_CAPS
    );
    
    // 读取配置空间数据（256字节）
    unsigned char config_data[256];  // 存储配置空间数据
    ssize_t nread = pread(
        device_fd,
        config_data,
        reg_info_config.size,
        reg_info_config.offset
    );

    if (nread != reg_info_config.size) {
        printf("Failed to read config space: %s (read %zd/%llu bytes)\n",
               strerror(errno), nread, (unsigned long long)reg_info_config.size);
        return -1;
    }

    // 解析配置空间关键信息（示例）
    unsigned short vendor_id = *(unsigned short*)(config_data + 0x00);  // 厂商ID（偏移0x00）
    unsigned short device_id = *(unsigned short*)(config_data + 0x02);  // 设备ID（偏移0x02）
    unsigned short command = *(unsigned short*)(config_data + 0x04);    // 命令寄存器（偏移0x04）
    unsigned short status = *(unsigned short*)(config_data + 0x06);     // 状态寄存器（偏移0x06）
    unsigned char class_code_base = config_data[0x0B];            // 基类别代码（偏移0x0B）
    unsigned char irq_pin = config_data[0x3D];                    // 中断引脚（偏移0x3D）

    // 打印关键信息
    printf("PCI Config Space Info:\n");
    printf("  Vendor ID: 0x%04X\n", vendor_id);
    printf("  Device ID: 0x%04X\n", device_id);
    printf("  Command Register: 0x%04X\n", command);
    printf("  Status Register: 0x%04X\n", status);
    printf("  Base Class Code: 0x%02X\n", class_code_base); //  (SATA通常为0x01，存储控制器)
    printf("  Interrupt Pin: %d (1=INTA#, 2=INTB#)\n", irq_pin);

    // 映射bar5空间
    struct vfio_region_info reg_info_bar5 = { .argsz = sizeof(reg_info_bar5), .index = VFIO_PCI_BAR5_REGION_INDEX };
    if (ioctl(device_fd, VFIO_DEVICE_GET_REGION_INFO, &reg_info_bar5)) {
        printf("Failed to get region info {%s}\n", strerror(errno));
        goto release;
    }
    printf("region_info index{%u} size{%lu} offset{%lu} READ{%d} WRITE{%d} MMAP{%d} CAPS{%d}\n",
        reg_info_bar5.index, reg_info_bar5.size, reg_info_bar5.offset,
        reg_info_bar5.flags & VFIO_REGION_INFO_FLAG_READ,
        reg_info_bar5.flags & VFIO_REGION_INFO_FLAG_WRITE,
        reg_info_bar5.flags & VFIO_REGION_INFO_FLAG_MMAP,
        reg_info_bar5.flags & VFIO_REGION_INFO_FLAG_CAPS
    );


    // 读取bar5数据（2048字节）
    unsigned char bar5_data[2048];  // 存储配置空间数据
    volatile uint8_t *bar5 = mmap(NULL, reg_info_bar5.size, PROT_READ | PROT_WRITE, MAP_SHARED, device_fd, reg_info_bar5.offset);
    if (bar5 == MAP_FAILED) {
        printf("Failed to map device memory {%s}.\n", strerror(errno));
        goto release;
    }
    volatile uint32_t *ghc = (uint32_t*)(bar5 + 0x04);
    *ghc |= (1 << 31);   // GHC.AE = 1
    *ghc |= (1 << 1);    // GHC.IE = 1
    volatile uint8_t *port = abar + 0x100 + 0x80 * i; // 端口i寄存器块

    // 10. 准备DMA缓冲区
    // void *dma_buffer = malloc(reg.size);
    // if (dma_buffer == NULL) {
    //     perror("Failed to malloc buffer.\n");
    //     goto release;
    // }
    // if (posix_memalign(&dma_buffer, sysconf(_SC_PAGESIZE), BUFFER_SIZE)) {
    //     perror("Failed to allocate the DMA buffer.\n");
    //     goto release;
    // }
    // // 读取厂商ID（偏移0x00，2字节）
    // uint16_t vendor_id = *(uint16_t *)(bar + 0x00);

    // // 读取设备ID（偏移0x02，2字节）
    // uint16_t device_id = *(uint16_t *)(bar + 0x02);

    // // 读取状态寄存器（偏移0x06，2字节）
    // uint16_t status = *(uint16_t *)(bar + 0x06);

    // // 读取命令寄存器（偏移0x04，2字节）
    // uint16_t command = *(uint16_t *)(bar + 0x04);

    // // 读取基址寄存器0 (BAR0，偏移0x10，4字节或8字节)
    // uint32_t bar0 = *(uint32_t *)(bar + 0x10);

    // // 打印结果
    // printf("Vendor ID: 0x%04x\n", vendor_id);
    // printf("Device ID: 0x%04x\n", device_id);
    // printf("Status: 0x%04x\n", status);
    // printf("Command: 0x%04x\n", command);
    // printf("BAR0: 0x%08x\n", bar0);
    // // 11. 模拟用户数据 (实际应用中从其他来源获取)
    // user_buffer = malloc(data_size);
    // if (!user_buffer) {
    //     perror("无法分配用户缓冲区");
    //     goto release;
    // }
    // memcpy(user_buffer, "Hello,World!", strlen("Hello,World!"));  // 填充测试数据
    
    // // 12. 拷贝数据到DMA缓冲区
    // memcpy(dma_buffer, user_buffer, data_size);
    
    // // 13. 将DMA缓冲区映射到IOMMU
    // struct vfio_iommu_type1_dma_map dma_map = {
    //     .argsz = sizeof(dma_map),
    //     .flags = VFIO_DMA_MAP_FLAG_READ | VFIO_DMA_MAP_FLAG_WRITE,
    //     .vaddr = (__u64)dma_buffer,
    //     .iova = 0,  // 由IOMMU分配
    //     .size = BUFFER_SIZE
    // };
    
    // if (ioctl(container_fd, VFIO_IOMMU_MAP_DMA, &dma_map)) {
    //     perror("无法映射DMA缓冲区");
    //     goto release;
    // }
    
    // printf("DMA缓冲区已设置，准备执行硬盘写入操作\n");
    
    // 注意：实际AHCI操作代码需要根据Intel AHCI规范实现
    // 这里只是框架示例
    
    // 14. 清理资源
    // munmap(bar0, reg_info_config.size);
    // free(dma_buffer);
    // free(user_buffer);
    // close(device_fd);
    // close(group_fd);
    // close(container_fd);
    // close(container_fd);
    
    // printf("操作完成\n");
    return 0;

release:
    if (group_fd >= 0) close(group_fd);
    if (container_fd >= 0) close(container_fd);
    if (device_fd >= 0) close(device_fd);
    return 0;
}

